#include "epsilon-greedy-allocation.h"
#include "ns3/log.h"
#include "ns3/simulator.h"
#include "ns3/mobility-model.h"
#include <algorithm>
#include <chrono>
#include <random>   // 添加头文件以支持随机数生成
#include <iomanip>  // 添加头文件以支持 std::setprecision

namespace ns3 {

NS_LOG_COMPONENT_DEFINE("EpsilonGreedyAllocation");

NS_OBJECT_ENSURE_REGISTERED(EpsilonGreedyAllocation);

TypeId EpsilonGreedyAllocation::GetTypeId()
{
    static TypeId tid = TypeId("ns3::EpsilonGreedyAllocation")
        .SetParent<ChannelAllocation>()
        .SetGroupName("SAGINIoV")
        .AddConstructor<EpsilonGreedyAllocation>();
    return tid;
}

EpsilonGreedyAllocation::EpsilonGreedyAllocation()
    : m_epsilon(0.1),
      m_throughputWeight(1.0),  // 添加吞吐量权重，默认为1.0
      m_delayWeight(0.1),       // 添加延迟权重，默认为0.1
      m_iteration(0)            // 初始化迭代计数器
{
    NS_LOG_FUNCTION(this);
}

EpsilonGreedyAllocation::~EpsilonGreedyAllocation()
{
    NS_LOG_FUNCTION(this);
}

void EpsilonGreedyAllocation::SetEpsilon(double epsilon)
{
    NS_LOG_FUNCTION(this << epsilon);
    m_epsilon = epsilon;
}

double EpsilonGreedyAllocation::GetEpsilon() const
{
    return m_epsilon;
}

// 性能指标计算函数中确保公平评估地面站和卫星性能
double EpsilonGreedyAllocation::CalculatePerformanceMetric(uint32_t nodeId, int channelId)
{
    // 计算SINR
    double sinr = CalculateSINR(nodeId, channelId);
    
    // 确定连接的节点类型
    NodeType connectedNodeType = GROUND_STATION;
    NodeContainer groundStations = m_nodeManager->GetGroundStationNodes();
    if (channelId >= static_cast<int>(groundStations.GetN())) {
        connectedNodeType = SATELLITE;
    }
    
    // 计算吞吐量和延迟
    double throughput = CalculateThroughput(sinr, connectedNodeType);
    double delay = CalculateDelay(nodeId, channelId);
    
    // 检查车辆到基站/卫星的距离
    double distance = 0.0;
    double elevationAngle = 0.0;
    Vector vehiclePos;
    
    // 获取车辆节点
    Ptr<Node> vehicleNode = nullptr;
    NodeContainer vehicles = m_nodeManager->GetVehicleNodes();
    for (uint32_t i = 0; i < vehicles.GetN(); i++) {
        if (vehicles.Get(i)->GetId() == nodeId) {
            vehicleNode = vehicles.Get(i);
            break;
        }
    }
    
    // 获取基站/卫星节点
    Ptr<Node> stationNode = nullptr;
    if (connectedNodeType == GROUND_STATION) {
        stationNode = groundStations.Get(channelId);
    } else {
        NodeContainer satellites = m_nodeManager->GetSatelliteNodes();
        int satIndex = channelId - groundStations.GetN();
        if (satIndex < static_cast<int>(satellites.GetN())) {
            stationNode = satellites.Get(satIndex);
        }
    }
    
    // 计算距离和仰角
    if (vehicleNode && stationNode) {
        Ptr<MobilityModel> vehicleMobility = vehicleNode->GetObject<MobilityModel>();
        Ptr<MobilityModel> stationMobility = stationNode->GetObject<MobilityModel>();
        
        if (vehicleMobility && stationMobility) {
            distance = vehicleMobility->GetDistanceFrom(stationMobility);
            
            if (connectedNodeType == SATELLITE) {
                vehiclePos = vehicleMobility->GetPosition();
                Vector satellitePos = stationMobility->GetPosition();
                double dx = satellitePos.x - vehiclePos.x;
                double dy = satellitePos.y - vehiclePos.y;
                double dz = satellitePos.z - vehiclePos.z;
                double horizontalDistance = std::sqrt(dx * dx + dy * dy);
                elevationAngle = std::atan2(dz, horizontalDistance) * 180.0 / M_PI;
            }
        }
    }
    
    // 检查连接是否可行
    bool connectionFeasible = false;
    if (connectedNodeType == GROUND_STATION) {
        connectionFeasible = IsGroundStationConnectionFeasible(distance, sinr);
    } else {
        connectionFeasible = IsSatelliteConnectionFeasible(elevationAngle, sinr);
    }
    
    // 如果连接不可行，返回极低的性能指标
    if (!connectionFeasible) {
        NS_LOG_DEBUG("节点" << nodeId << "到" << (connectedNodeType == GROUND_STATION ? "地面站" : "卫星") 
                  << "(" << channelId << ")的连接不可行，性能指标设为0");
        return 0.0;
    }
    
    // 性能指标计算 - 确保地面站和卫星公平竞争
    double throughputWeight = m_throughputWeight; 
    double delayWeight = m_delayWeight;
    
    // 根据节点类型动态调整权重，避免偏好
    if (connectedNodeType == GROUND_STATION) {
        // 稍微提高地面站的延迟权重，因为地面站通常有更低的延迟
        throughputWeight *= 1.0;
        delayWeight *= 1.1;
    } else {
        // 稍微提高卫星的吞吐量权重，因为卫星通常有更高的吞吐量
        throughputWeight *= 1.1;
        delayWeight *= 1.0;
    }
    
    // 正常计算性能指标
    double metric = throughputWeight * throughput - delayWeight * (delay * 1000); // 转换为毫秒
    
    NS_LOG_DEBUG("性能指标计算: 节点=" << nodeId << ", 信道=" << channelId 
               << ", 类型=" << (connectedNodeType == GROUND_STATION ? "地面站" : "卫星")
               << ", 距离=" << distance << "m"
               << (connectedNodeType == SATELLITE ? ", 仰角=" + std::to_string(elevationAngle) + "°" : "")
               << ", SINR=" << sinr << "dB"
               << ", 吞吐量=" << throughput << "Mbps"
               << ", 延迟=" << (delay * 1000) << "ms"
               << ", 性能指标=" << metric);
    
    return metric;
}

// 探索阶段随机选择信道时，确保随机性真实无偏
int EpsilonGreedyAllocation::ExploreRandomChannel(uint32_t nodeId, const std::vector<int>& availableChannels)
{
    NS_LOG_FUNCTION(this << nodeId);
    
    if (availableChannels.empty()) {
        NS_LOG_WARN("没有可用的信道");
        return -1;
    }
    
    // 使用真正均匀的随机分布
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, availableChannels.size() - 1);
    
    int randomIndex = dis(gen);
    int selectedChannel = availableChannels[randomIndex];
    
    NS_LOG_DEBUG("节点 " << nodeId << " 探索随机选择信道: " << selectedChannel);
    
    return selectedChannel;
}

std::map<uint32_t, int> EpsilonGreedyAllocation::AllocateChannels()
{
    NS_LOG_FUNCTION(this);
    
    if (!m_nodeManager) {
        NS_LOG_ERROR("Node manager is not set");
        return std::map<uint32_t, int>();
    }
    
    // 获取节点
    NodeContainer vehicles = m_nodeManager->GetVehicleNodes();
    NodeContainer groundStations = m_nodeManager->GetGroundStationNodes();
    NodeContainer satellites = m_nodeManager->GetSatelliteNodes();
    
    int numGroundStations = groundStations.GetN();
    int numSatellites = satellites.GetN();
    
    // 获取当前时间
    double currentTime = Simulator::Now().GetSeconds();
    
    // 获取当前活动的车辆数据
    std::vector<VehicleData> vehicleData = m_nodeManager->GetVehicleDataAtTime(currentTime);
    NS_LOG_INFO("Found " << vehicleData.size() << " active vehicles at time " << currentTime << "s");
    
    if (vehicleData.empty()) {
        NS_LOG_INFO("No active vehicles at current time, nothing to allocate");
        return std::map<uint32_t, int>();
    }
    
    // 获取可用信道
    std::vector<int> availableChannels;
    
    // 结果保存
    std::map<uint32_t, int> allocation;
    std::map<uint32_t, double> bestMetrics;
    
    // 随机种子使用当前时间和迭代次数来确保每次不同
    std::srand(static_cast<unsigned int>(Simulator::Now().GetMilliSeconds() + m_iteration++));
    
    // 为每个活跃的车辆分配信道
    for (const auto& vehicle : vehicleData) {
        // 找到对应的节点ID
        uint32_t vehicleNodeId = 0;
        
        // 遍历所有车辆节点，查找匹配vehicleId的节点
        for (uint32_t i = 0; i < vehicles.GetN(); i++) {
            std::string id = m_nodeManager->GetVehicleIdForNode(vehicles.Get(i)->GetId());
            if (id == vehicle.id) {
                vehicleNodeId = vehicles.Get(i)->GetId();
                break;
            }
        }
        
        if (vehicleNodeId == 0) {
            NS_LOG_WARN("Could not find node for vehicle ID: " << vehicle.id);
            continue; // 无效的节点ID
        }
        
        // 获取车辆节点
        Ptr<Node> vehicleNode = nullptr;
        for (uint32_t i = 0; i < vehicles.GetN(); i++) {
            if (vehicles.Get(i)->GetId() == vehicleNodeId) {
                vehicleNode = vehicles.Get(i);
                break;
            }
        }
        
        if (!vehicleNode) {
            continue;
        }
        
        // 获取车辆移动模型
        Ptr<MobilityModel> vehicleMobility = vehicleNode->GetObject<MobilityModel>();
        if (!vehicleMobility) {
            continue;
        }
        
        // 创建可用信道列表
        availableChannels.clear();
        
        // 添加地面站信道 (如果允许)
        if (m_allowGroundStationConnection) {
            for (int i = 0; i < numGroundStations; i++) {
                // 获取地面站节点
                Ptr<Node> stationNode = groundStations.Get(i);
                Ptr<MobilityModel> stationMobility = stationNode->GetObject<MobilityModel>();
                
                if (stationMobility) {
                    // 计算距离
                    double distance = vehicleMobility->GetDistanceFrom(stationMobility);
                    
                    // 计算SINR
                    double sinr = CalculateSINR(vehicleNodeId, i);
                    
                    // 检查连接是否可行
                    if (IsGroundStationConnectionFeasible(distance, sinr)) {
                        availableChannels.push_back(i);
                    }
                }
            }
        }
        
        // 添加卫星信道 (如果允许)
        if (m_allowSatelliteConnection) {
            for (int i = 0; i < numSatellites; i++) {
                // 获取卫星节点
                Ptr<Node> satelliteNode = satellites.Get(i);
                Ptr<MobilityModel> satelliteMobility = satelliteNode->GetObject<MobilityModel>();
                
                if (satelliteMobility) {
                    // 计算卫星仰角
                    Vector vehiclePos = vehicleMobility->GetPosition();
                    Vector satellitePos = satelliteMobility->GetPosition();
                    double dx = satellitePos.x - vehiclePos.x;
                    double dy = satellitePos.y - vehiclePos.y;
                    double dz = satellitePos.z - vehiclePos.z;
                    double horizontalDistance = std::sqrt(dx * dx + dy * dy);
                    double elevationAngle = std::atan2(dz, horizontalDistance) * 180.0 / M_PI;
                    
                    // 计算SINR
                    double sinr = CalculateSINR(vehicleNodeId, numGroundStations + i);
                    
                    // 检查连接是否可行
                    if (IsSatelliteConnectionFeasible(elevationAngle, sinr)) {
                        availableChannels.push_back(numGroundStations + i);
                    }
                }
            }
        }
        
        if (availableChannels.empty()) {
            NS_LOG_INFO("No available channels for vehicle " << vehicle.id 
                        << " (node " << vehicleNodeId << ")");
            continue;
        }
        
        // 以epsilon的概率进行探索(随机选择信道)
        bool explore = (static_cast<double>(std::rand()) / RAND_MAX) < m_epsilon;
        
        if (explore) {
            // 确保随机信道选择不偏向地面站或卫星
            int randomIndex = std::rand() % availableChannels.size();
            int selectedChannel = availableChannels[randomIndex];
            
            allocation[vehicleNodeId] = selectedChannel;
            
            // 计算并记录所选信道的性能指标
            double metric = CalculatePerformanceMetric(vehicleNodeId, selectedChannel);
            bestMetrics[vehicleNodeId] = metric;
            
            NS_LOG_DEBUG("节点 " << vehicleNodeId << " (车辆 " << vehicle.id << ") 探索阶段选择信道 " << selectedChannel 
                      << ", 性能指标: " << metric);
        } else {
            // 开发阶段：选择最佳信道
            double bestMetric = -std::numeric_limits<double>::infinity();
            int bestChannel = -1;
            
            // 测试所有可用信道
            for (int channelId : availableChannels) {
                // 计算性能指标
                double metric = CalculatePerformanceMetric(vehicleNodeId, channelId);
                
                // 确定节点类型，用于记录
                NodeType connectedNodeType = GROUND_STATION;
                if (channelId >= numGroundStations) {
                    connectedNodeType = SATELLITE;
                }
                
                // 更新最佳指标和信道
                if (metric > bestMetric) {
                    bestMetric = metric;
                    bestChannel = channelId;
                    
                    NS_LOG_DEBUG("节点 " << vehicleNodeId << " (车辆 " << vehicle.id << ") 更新最佳信道为 " << channelId 
                              << " (" << (connectedNodeType == GROUND_STATION ? "地面站" : "卫星") << ")"
                              << ", 性能指标: " << metric);
                }
            }
            
            if (bestChannel >= 0) {
                allocation[vehicleNodeId] = bestChannel;
                bestMetrics[vehicleNodeId] = bestMetric;
                
                NS_LOG_DEBUG("节点 " << vehicleNodeId << " (车辆 " << vehicle.id << ") 最终选择信道 " << bestChannel 
                          << ", 性能指标: " << bestMetric);
            }
        }
    }
    
    // 记录最终指标和选择结果
    NS_LOG_INFO("Epsilon-greedy分配结果，时间 " << currentTime << "s (epsilon=" << m_epsilon << "):");
    int groundStationCount = 0;
    int satelliteCount = 0;
    
    for (const auto& [nodeId, channelId] : allocation) {
        bool isGroundStation = channelId < numGroundStations;
        std::string vehicleId = m_nodeManager->GetVehicleIdForNode(nodeId);
        
        if (isGroundStation) {
            groundStationCount++;
            NS_LOG_INFO("Vehicle " << vehicleId << " (node " << nodeId 
                       << ") allocated to ground station " << channelId
                       << ", 性能指标: " << bestMetrics[nodeId]);
        } else {
            satelliteCount++;
            NS_LOG_INFO("Vehicle " << vehicleId << " (node " << nodeId 
                       << ") allocated to satellite " << (channelId - numGroundStations)
                       << ", 性能指标: " << bestMetrics[nodeId]);
        }
    }
    
    NS_LOG_INFO("总分配: " << allocation.size() << " 车辆, "
              << groundStationCount << " 连接地面站, "
              << satelliteCount << " 连接卫星");
    
    return allocation;
}

void EpsilonGreedyAllocation::LogAllocationResults()
{
    NS_LOG_FUNCTION(this);
    
    if (!m_nodeManager) {
        NS_LOG_ERROR("Node manager is not set");
        return;
    }
    
    // 获取当前时间
    double currentTime = Simulator::Now().GetSeconds();
    
    // 获取信道分配
    std::map<uint32_t, int> allocation = GetCurrentAllocation();
    
    // 获取节点
    NodeContainer groundStations = m_nodeManager->GetGroundStationNodes();
    NodeContainer satellites = m_nodeManager->GetSatelliteNodes();
    
    int numGroundStations = groundStations.GetN();
    
    // 日志输出当前分配情况
    NS_LOG_INFO("Current channel allocation at time " << currentTime << "s:");
    
    for (const auto& [nodeId, channelId] : allocation) {
        std::string vehicleId = m_nodeManager->GetVehicleIdForNode(nodeId);
        
        if (channelId < numGroundStations) {
            NS_LOG_INFO("Vehicle " << vehicleId << " (node " << nodeId 
                       << ") is connected to ground station " << channelId);
        } else {
            NS_LOG_INFO("Vehicle " << vehicleId << " (node " << nodeId 
                       << ") is connected to satellite " << (channelId - numGroundStations));
        }
    }
}

}  // namespace ns3
